import sys
from PySide6.QtWidgets import (QApplication, QMainWindow, QLabel, QScrollArea, QVBoxLayout, QWidget, QFileDialog)
from PySide6.QtGui import (QPixmap, QImage, QPainter, QPen, QColor, QAction, Qt)
from PySide6.QtCore import QPoint, Slot


class AnnotatableImage(QLabel):
    def __init__(self, pixmap):
        super().__init__()
        self.setPixmap(pixmap)
        self.pixmap = pixmap.copy()  # 创建原始图像的副本以进行绘制
        self.drawing = False
        self.lastPoint = QPoint()
        self.penColor = QColor(Qt.red)  # 默认绘图颜色为红色
        self.penWidth = 2  # 默认线条宽度

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.drawing = True
            self.lastPoint = event.position().toPoint()

    def mouseMoveEvent(self, event):
        if (event.buttons() & Qt.LeftButton) and self.drawing:
            painter = QPainter(self.pixmap)
            pen = QPen(self.penColor, self.penWidth, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin)
            painter.setPen(pen)
            painter.drawLine(self.lastPoint, event.position().toPoint())
            self.lastPoint = event.position().toPoint()
            self.setPixmap(self.pixmap)
            self.update()

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.drawing = False

class ImageWindow(QMainWindow):
    def __init__(self, image_path):
        super().__init__()

        self.pixmap = QPixmap(image_path)
        self.image_label = AnnotatableImage(self.pixmap)

        self.scroll_area = QScrollArea()
        # self.scroll_area.setBackgroundRole(Qt.LightBackground)
        self.scroll_area.setWidgetResizable(True)
        self.scroll_area.setWidget(self.image_label)

        self.main_widget = QWidget()
        self.layout = QVBoxLayout(self.main_widget)
        self.layout.addWidget(self.scroll_area)
        self.setCentralWidget(self.main_widget)

        # 添加菜单栏选项用于保存带注释的图片
        save_action = QAction("Save", self)
        save_action.triggered.connect(self.save_annotated_image)
        menu_bar = self.menuBar()
        file_menu = menu_bar.addMenu("File")
        file_menu.addAction(save_action)

        self.resizeEvent = self.on_resize

    @Slot()
    def on_resize(self, event):
        self.update_image_size()

    def update_image_size(self):
        new_width = self.width()
        scaled_pixmap = self.pixmap.scaledToWidth(new_width, Qt.SmoothTransformation)
        self.image_label.pixmap = scaled_pixmap
        self.image_label.setPixmap(scaled_pixmap)
        self.image_label.adjustSize()

    def save_annotated_image(self):
        options = QFileDialog.Options()
        file_name, _ = QFileDialog.getSaveFileName(self, "Save Annotated Image", "", "Images (*.png *.xpm *.jpg);;All Files (*)", options=options)
        if file_name:
            self.image_label.pixmap.save(file_name)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    image_path = "../asset/gammaImg.png"  # 替换为你的图片路径
    window = ImageWindow(image_path)
    window.setWindowTitle("Annotate Image Viewer")
    window.show()
    sys.exit(app.exec())